# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np
import const
import random

class DataParse(object):
	def __init__(self):
		self.dataDir = const.DATA_DIR
		self.dataFilesName = const.GESTURES
		self.dataFilesType = const.FILE_TYPE
		self.fileNameToLabel = const.GESTURES_TO_Y

		self.numOfSample = 0
		self.timeStepsEachSample = const.T * const.AUG_STEPS
		self.variablesEachSample = const.V
		self.classicNumOfSample = const.TYPE_NUM

		self.xData = []
		self.yData = []
		self.shapeOfX = []
		self.shapeOfY = []

		self._initData()

	def _initData(self):
		for name in self.dataFilesName:
			try:
				df = pd.read_csv(self.dataDir + name + self.dataFilesType)
				dfList = df.values.tolist()
				timeStepsNum = df.shape[0]
				variablesNum = df.shape[1]
			except Exception as e:
				print ('ERROR in read %s, %s'%(name, e))
				continue

			assert(timeStepsNum % self.timeStepsEachSample == 0 or variablesNum == self.variablesEachSample)
			assert(len(self.fileNameToLabel[name]) == self.classicNumOfSample)

			sampleNumInThisFile = int(timeStepsNum/self.timeStepsEachSample)
			self.numOfSample += sampleNumInThisFile
			self.xData += [dfList[i:i+self.timeStepsEachSample] for i in range(0, len(dfList), self.timeStepsEachSample)]
			self.yData += [self.fileNameToLabel[name]] * sampleNumInThisFile

		self.shapeOfX = [self.numOfSample, self.timeStepsEachSample, self.variablesEachSample]
		self.shapeOfY = [self.numOfSample, self.classicNumOfSample]
	
	def getData(self, returnType=const.READ_DATA_ARRAY_TYPE):
		if returnType == const.READ_DATA_LIST_TYPE:
			return self.xData, self.yData
		if returnType == const.READ_DATA_ARRAY_TYPE:
			return np.array(self.xData), np.array(self.yData)

	def generateTrainAndTestData(self, returnType=const.READ_DATA_ARRAY_TYPE, ratio=0.9):
		combined = list(zip(self.xData, self.yData))
		random.shuffle(combined)
		x, y = zip(*combined)

		trainNum = int(self.numOfSample * ratio)
		testNum = self.numOfSample - trainNum
		trainX = x[:trainNum]
		trainY = y[:trainNum]
		testX = x[trainNum:]
		testY = y[trainNum:]
		if returnType == const.READ_DATA_ARRAY_TYPE:
			return trainNum, np.array(trainX), np.array(trainY), testNum, np.array(testX), np.array(testY)
		if returnType == const.READ_DATA_LIST_TYPE:
			return trainNum, trainX, trainY, testNum, testX, testY
	
	def dataAugmentation(self, x, y):
		if x.shape[0] != y.shape[0]:
			print ('augmentation error')
			exit(0)
		xAug = []
		yAug = []
		for i in range(len(x)):
			for j in range(0, const.AUG_STEPS):
				xAug.append([x[i][k].tolist() for k in range(j, len(x[i]), const.AUG_STEPS)])
				yAug.append(y[i])
		return np.array(xAug), np.array(yAug)